package com.wl.mr.output;

import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.io.IOUtils;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;

import java.io.IOException;

public class CusterRecordWriter extends RecordWriter<Text, NullWritable> {

    private FSDataOutputStream lagou;
    private FSDataOutputStream other;

    public CusterRecordWriter(FSDataOutputStream lagou, FSDataOutputStream other) {
        this.lagou = lagou;
        this.other = other;
    }

    @Override
    public void write(Text key, NullWritable value) throws IOException, InterruptedException {
        if(key.toString().contains("lagou")){
            lagou.write(key.toString().getBytes());
            lagou.write("\r\n".getBytes());
        }else {
            other.write(key.toString().getBytes());
            other.write("\r\n".getBytes());
        }
    }

    @Override
    public void close(TaskAttemptContext context) throws IOException, InterruptedException {
        IOUtils.closeStream(lagou);
        IOUtils.closeStream(other);
    }
}
